# Time-stamp: <2008-10-02 11:11:02 Tao Liu>

"""Module Description

Copyright (c) 2008 Yong Zhang, Tao Liu <taoliu@jimmy.harvard.edu>

This code is free software; you can redistribute it and/or modify it
under the terms of the Artistic License (see the file COPYING included
with the distribution).

@status:  experimental
@version: $Revision$
@author:  Yong Zhang, Tao Liu
@contact: taoliu@jimmy.harvard.edu
"""
from math import log as mathlog
from MACS.OutputWriter import zwig_write
from MACS.IO import PeakIO
from MACS.Prob import poisson_cdf,poisson_cdf_inv
from MACS.Constants import *


class PeakDetect:
    """Class to do the peak calling.

    e.g:
    >>> from MACS.PeakDetect import PeakDetect
    >>> pd = PeakDetect(treat=treatdata, control=controldata, pvalue=pvalue_cutoff, d=100, scan_window=200, gsize=3000000000)
    >>> pd.call_peaks()
    >>> print pd.toxls()
    """
    def __init__ (self,opt=None,treat=None, control=None):
        """Initialize the PeakDetect object.

        """
        self.opt  = opt
        self.info = opt.info
        self.debug = opt.debug
        self.warn = opt.warn
        self.treat = treat
        self.control = control
        self.ratio_treat2control = None
        self.peaks = None
        self.final_peaks = None
        self.final_negative_peaks = None
        
        self.pvalue = opt.log_pvalue
        self.d = opt.d
        self.shift_size = self.d/2
        self.scan_window = opt.scanwindow
        self.gsize = opt.gsize
        
        self.nolambda = opt.nolambda
        (self.first_lambda_region,self.second_lambda_region,self.third_lambda_region) = opt.lambdaset
        if (self.nolambda):
            self.info("#3 !!!! DYNAMIC LAMBDA IS DISABLED !!!!")
        self.diag = opt.diag
        self.save_wig = opt.store_wig
        self.zwig_tr = opt.zwig_tr
        self.zwig_ctl= opt.zwig_ctl

    def call_peaks (self):
        """Call peaks function.

        Scan the whole genome for peaks. RESULTS WILL BE SAVED IN
        self.final_peaks and self.final_negative_peaks.
        """
        if self.control:
            self.peaks = self.__call_peaks_w_control ()
        else:
            self.peaks = self.__call_peaks_wo_control ()
        return None

    def diag_result (self):
        """Run the diagnosis process on sequencing saturation.
        
        """
        if not self.diag:
            return None
        if self.control:
            return self.__diag_w_control()
        else:
            return self.__diag_wo_control()

    def toxls (self):
        """
        """
        text = ""
        if self.control and self.peaks:
            text += "\t".join(("chr","start", "end",  "length",  "summit", "tags", "-10*log10(pvalue)", "fold_enrichment", "FDR(%)"))+"\n"
        elif self.peaks:
            text += "\t".join(("chr","start", "end",  "length",  "summit", "tags", "-10*log10(pvalue)", "fold_enrichment"))+"\n"
        else:
            return None
        
        chrs = self.peaks.keys()
        chrs.sort()
        for chrom in chrs:
            for peak in self.peaks[chrom]:
                text += "%s\t%d\t%d\t%d" % (chrom,peak[0]+1,peak[1],peak[2])
                peak_summit_relative_pos = peak[3]-peak[0]+1
                text += "\t%d" % (peak_summit_relative_pos)
                text += "\t%d\t%.2f" % (peak[5],peak[6])
                text += "\t%.2f" % (peak[7])
                if self.control:
                    if peak[8]>=100:
                        text += "\t100"
                    else:
                        text += "\t%.2f" % (peak[8])
                text+= "\n"
        return text

    def neg_toxls (self):
        text = ""
        text += "\t".join(("chr","start", "end",  "length",  "summit", "tags", "-10*log10(pvalue)","fold_enrichment"))+"\n"
        chrs = self.final_negative_peaks.keys()
        chrs.sort()
        for chrom in chrs:
            for peak in self.final_negative_peaks[chrom]:
                text += "%s\t%d\t%d\t%d" % (chrom,peak[0]+1,peak[1],peak[2])
                peak_summit_relative_pos = peak[3]-peak[0]+1
                text += "\t%d" % (peak_summit_relative_pos)
                text += "\t%d\t%.2f" % (peak[5],peak[6])
                text += "\t%.2f" % (peak[7])
                text+= "\n"
        return text

    def tobed (self):
        text = ""
        chrs = self.peaks.keys()
        chrs.sort()
        for chrom in chrs:
            for peak in self.peaks[chrom]:
                text+= "%s\t%d\t%d\n" % (chrom,peak[0],peak[1])
        return text

    def __add_fdr (self, final, negative): 
        """
        A peak info type is a: dictionary

        key value: chromosome

        items: (peak start,peak end, peak length, peak summit, peak
        height, number of tags in peak region, peak pvalue, peak
        fold_enrichment, fdr) <-- tuple type
        """
        pvalue2fdr = {}
        pvalues_final = []
        pvalues_negative = []
        chrs = final.keys()
        a = pvalues_final.append
        for chrom in chrs:
            for i in final[chrom]:
                a(i[6]) # i[6] is pvalue in peak info
                pvalue2fdr[i[6]]=None
        chrs = negative.keys()
        a = pvalues_negative.append
        for chrom in chrs:
            for i in negative[chrom]:
                a(i[6])
        pvalues_final.sort(reverse=True)
        pvalues_final_l = len(pvalues_final)
        pvalues_negative.sort(reverse=True)
        pvalues_negative_l = len(pvalues_negative)        
        pvalues = pvalue2fdr.keys()
        pvalues.sort(reverse=True)
        index_p2f_pos = 0
        index_p2f_neg = 0
        for p in pvalues:
            while index_p2f_pos<pvalues_final_l and p<=pvalues_final[index_p2f_pos]:
                index_p2f_pos += 1
            n_final = index_p2f_pos

            while  index_p2f_neg<pvalues_negative_l and p<=pvalues_negative[index_p2f_neg]:
                index_p2f_neg += 1
            n_negative = index_p2f_neg
            pvalue2fdr[p] = 100.0 * n_negative / n_final

        new_info = {}
        chrs = final.keys()
        for chrom in chrs:
            new_info[chrom] = []
            for i in final[chrom]:
                tmp = list(i)
                tmp.append(pvalue2fdr[i[6]])
                new_info[chrom].append(tuple(tmp))      # i[6] is pvalue in peak info
        return new_info

    def __call_peaks_w_control (self):
        """To call peaks with control data.

        A peak info type is a: dictionary

        key value: chromosome

        items: (peak start,peak end, peak length, peak summit, peak
        height, number of tags in peak region, peak pvalue, peak
        fold_enrichment) <-- tuple type
        """
        self.lambda_bg = float(self.scan_window)*self.treat.total_unique/self.gsize
        self.debug("#3 background lambda: %.2f " % (self.lambda_bg))
        self.min_tags = poisson_cdf_inv(1-pow(10,self.pvalue/-10),self.lambda_bg)+1
        self.debug("#3 min tags: %d" % (self.min_tags))

        self.ratio_treat2control = float(self.treat.total_unique)/self.control.total_unique
        if self.ratio_treat2control > 2 or self.ratio_treat2control < 0.5:
            self.warn("Treatment tags and Control tags are uneven! FDR may be wrong!")
        self.info("#3 shift treatment data")
        self.__shift_trackI(self.treat)
        self.info("#3 merge +/- strand of treatment data")
        self.treat.merge_plus_minus_ranges_w_duplicates ()
        self.debug("#3 after shift and merging, tags: %d" % (self.treat.total))
        if self.save_wig:
            self.info("#3 save the shifted and merged tag counts into wiggle file...")
            zwig_write(self.treat,self.opt.wig_dir_tr,self.zwig_tr,self.d,log=self.info)
        self.info("#3 call peak candidates")
        peak_candidates = self.__call_peaks_from_trackI (self.treat)

        self.info("#3 shift control data")
        self.info("#3 merge +/- strand of control data")
        self.__shift_trackI(self.control)
        self.control.merge_plus_minus_ranges_w_duplicates ()
        self.debug("#3 after shift and merging, tags: %d" % (self.control.total))
        if self.save_wig:
            self.info("#3 save the shifted and merged tag counts into wiggle file...")
            zwig_write(self.control,self.opt.wig_dir_ctl,self.zwig_ctl,self.d,log=self.info)
        self.info("#3 call negative peak candidates")
        negative_peak_candidates = self.__call_peaks_from_trackI (self.control)
        
        self.info("#3 use control data to filter peak candidates...")
        self.final_peaks = self.__filter_w_control(peak_candidates,self.control, self.ratio_treat2control,fake_when_missing_control=True)
        self.info("#3 find negative peaks by reversing treat and control")
        self.final_negative_peaks = self.__filter_w_control(negative_peak_candidates,self.treat, 1.0/self.ratio_treat2control)
        return self.__add_fdr (self.final_peaks, self.final_negative_peaks)

    def __call_peaks_wo_control (self):
        """To call peaks w/o control data.

        """
        self.lambda_bg = float(self.scan_window)*self.treat.total_unique/self.gsize
        self.debug("#3 background lambda: %.2f " % (self.lambda_bg))
        self.min_tags = poisson_cdf_inv(1-pow(10,self.pvalue/-10),self.lambda_bg)+1
        self.debug("#3 min tags: %d" % (self.min_tags))

        self.info("#3 shift treatment data")
        self.__shift_trackI(self.treat)
        self.info("#3 merge +/- strand of treatment data")
        self.treat.merge_plus_minus_ranges_w_duplicates ()
        self.debug("#3 after shift and merging, tags: %d" % (self.treat.total))
        if self.save_wig:
            self.info("#3 save the shifted and merged tag counts into wiggle file...")
            zwig_write(self.treat,self.opt.wig_dir_tr,self.zwig_tr,self.d,log=self.info)
        self.info("#3 call peak candidates")
        peak_candidates = self.__call_peaks_from_trackI (self.treat)
        self.info("#3 use self to calculate local lambda and  filter peak candidates...")
        self.final_peaks = self.__filter_w_control(peak_candidates,self.treat,1,pass_1k=True)
        return self.final_peaks

    def __print_peak_info (self, peak_info):
        """Print out peak information.

        A peak info type is a: dictionary

        key value: chromosome

        items: (peak start,peak end, peak length, peak summit, peak
        height, number of tags in peak region, peak pvalue, peak
        fold_enrichment) <-- tuple type
        """
        chrs = peak_info.keys()
        chrs.sort()
        for chrom in chrs:
            peak_list = peak_info[chrom]
            for peak in peak_list:
                print chrom+"\t"+"\t".join(map(str,peak))

    def __filter_w_control (self, peak_info, control, treat2control_ratio, pass_1k=False, write2wig= False, fake_when_missing_control=False ):
        """Use control data to calculate several lambda values around
        1k, 5k and 10k region around peak summit. Choose the highest
        one as local lambda, then calculate p-value in poisson
        distribution.

        Return value type in this format:
        a dictionary
        key value : chromosome
        items : array of (peak_start,peak_end,peak_length,peak_summit,peak_height,peak_num_tags,peak_pvalue,peak_fold_enrichment)
        """
        final_peak_info = {}
        chrs = peak_info.keys()
        chrs.sort()
        total = 0
        for chrom in chrs:
            self.debug("#3 Chromosome %s" % (chrom))
            n_chrom = 0
            final_peak_info[chrom] = []
            peak_list = peak_info[chrom]
            try:
                (ctags,tmp) = control.get_ranges_by_chr(chrom)
            except:
                self.warn("Missing %s data, skip it..." % (chrom))
                if fake_when_missing_control:
                    ctags = [-1,]
                    self.warn("Fake a tag at %s:%d" % (chrom,-1))
                    tmp=[]
                else:
                    continue
            index_ctag = 0      # index for control tags
            flag_find_ctag_in_peak = False
            prev_index_ctag = 0
            len_ctags =len(ctags)
            for i in range(len(peak_list)):
                (peak_start,peak_end,peak_length,peak_summit,peak_height,peak_num_tags) = peak_list[i]

                window_size_4_lambda = max(peak_length,self.scan_window)
                lambda_bg = self.lambda_bg/self.scan_window*window_size_4_lambda                
                # no lambda?
                if self.nolambda:
                    local_lambda = lambda_bg
                else:
                    left_10k = peak_summit-self.third_lambda_region/2
                    left_5k = peak_summit-self.second_lambda_region/2
                    left_1k = peak_summit-self.first_lambda_region/2
                    right_10k = peak_summit+self.third_lambda_region/2
                    right_5k = peak_summit+self.second_lambda_region/2
                    right_1k = peak_summit+self.first_lambda_region/2
                    num_10k = 0
                    num_5k = 0
                    num_1k = 0
                    while index_ctag < len_ctags:
                        if ctags[index_ctag] < left_10k:
                            # go to next control tag
                            index_ctag+=1
                        elif right_10k < ctags[index_ctag]:
                            # finalize and go to next peak region
                            flag_find_ctag_in_peak = False
                            index_ctag = prev_index_ctag 
                            break
                        else:
                            # in -10k ~ 10k region
                            if not flag_find_ctag_in_peak:
                                flag_find_ctag_in_peak = True
                                prev_index_ctag = index_ctag
                            num_10k +=1
                            if left_1k <ctags[index_ctag]< right_1k:
                                num_1k+=1
                                num_5k+=1
                            elif left_5k <ctags[index_ctag]< right_5k:
                                num_5k+=1
                            index_ctag += 1

                    lambda_10k = float(num_10k)/self.third_lambda_region*treat2control_ratio*window_size_4_lambda
                    lambda_5k = float(num_5k)/self.second_lambda_region*treat2control_ratio*window_size_4_lambda
                    lambda_1k = float(num_1k)/self.first_lambda_region*treat2control_ratio*window_size_4_lambda

                    if pass_1k:
                        local_lambda = max(lambda_bg,lambda_10k,lambda_5k)
                    else:
                        local_lambda = max(lambda_bg,lambda_10k,lambda_5k,lambda_1k)
                #print local_lambda,peak_start,peak_end
                p_tmp = poisson_cdf(peak_num_tags,local_lambda,lower=False)
                if p_tmp <= 0:
                    peak_pvalue = 3100
                else:
                    peak_pvalue = mathlog(p_tmp,10) * -10
                #print local_lambda,peak_pvalue
                if peak_pvalue > self.pvalue:
                    n_chrom += 1
                    total += 1
                    peak_fold_enrichment = float(peak_height)/local_lambda*window_size_4_lambda/self.d
                    final_peak_info[chrom].append((peak_start,peak_end,peak_length,peak_summit,peak_height,peak_num_tags,peak_pvalue,peak_fold_enrichment))

            self.debug("#3 peaks whose pvalue < cutoff: %d" % (n_chrom))
        self.info("#3 Finally, %d peaks are called!" % (total))
        return final_peak_info

    def __call_peaks_from_trackI (self, trackI):
        """Call peak candidates from trackI data. Using every tag as
        step and scan the self.scan_window region around the tag. If
        tag number is greater than self.min_tags, then the position is
        recorded.

        Return: data in this format. (peak_start,peak_end,peak_length,peak_summit,peak_height,peak_num_tags)
        """
        peak_candidates = {}
        self.debug("#3 search peak condidates...")
        chrs = trackI.get_chr_names()
        total = 0
        for chrom in chrs:
            self.debug("#3 Chromosome %s" % (chrom))
            n_chrom = 0
            peak_candidates[chrom] = []
            (tags,tmp) = trackI.get_ranges_by_chr(chrom)
            len_t = len(tags)
            cpr_tags = []       # Candidate Peak Region tags
            cpr_tags.extend(tags[:self.min_tags-1])
            number_cpr_tags = self.min_tags-1
            p = self.min_tags-1 # Next Tag Index
            while p < len_t:
                if number_cpr_tags >= self.min_tags:
                    if tags[p] - cpr_tags[-1*self.min_tags+1] <= self.scan_window:
                        # add next tag, if the new tag is less than self.scan_window away from previous no. self.min_tags tag
                        cpr_tags.append(tags[p])
                        number_cpr_tags += 1
                        p+=1
                    else:
                        # candidate peak region is ready, call peak...
                        (peak_start,peak_end,peak_length,peak_summit,peak_height) = self.__tags_call_peak (cpr_tags)
                        peak_candidates[chrom].append((peak_start,peak_end,peak_length,peak_summit,peak_height,number_cpr_tags))
                        cpr_tags = [tags[p]] # reset

                        number_cpr_tags = 1
                        total += 1
                        n_chrom += 1
                else:
                    # add next tag, but if the first one in cpr_tags
                    # is more than self.scan_window away from this new
                    # tag, remove the first one from the list
                    if tags[p] - cpr_tags[0] >= self.scan_window:
                        cpr_tags.pop(0)
                        number_cpr_tags -= 1
                    cpr_tags.append(tags[p])
                    number_cpr_tags += 1
                    p+=1
            self.debug("#3 peak candidates: %d" % (n_chrom))
        self.debug("#3 Total number of candidates: %d" % (total))
        return self.__remove_overlapping_peaks(peak_candidates)
                
    def __tags_call_peak (self, tags ):
        """Project tags to a line. Then find the highest point.

        """
        start = tags[0]-self.d/2
        end = tags[-1]+self.d/2+1
        region_length = end - start
        line= [0]*region_length
        for tag in tags:
            tag_projected_start = tag-start-self.d/2
            tag_projected_end = tag-start+self.d/2
            for i in range(tag_projected_start,tag_projected_end):
                line[i] += 1
        tops = []
        top_height = 0
        for i in range(len(line)):
            if line[i] > top_height:
                top_height = line[i]
                tops = [i]
            elif line[i] == top_height:
                tops.append(i)
        peak_summit = tops[len(tops)/2]+start
        return (start,end,region_length,peak_summit,top_height)
            

    def __shift_trackI (self, trackI):
        """Shift trackI data to right (for plus strand) or left (for
        minus strand).

        """
        chrs = trackI.get_chr_names()
        for chrom in chrs:
            tags = trackI.get_ranges_by_chr(chrom)
            # plus
            for i in range(len(tags[0])):
                tags[0][i]+=self.shift_size
            # minus
            for i in range(len(tags[1])):
                tags[1][i]-=self.shift_size
        return
    

    def __diag_w_control (self):
        # sample
        sample_peaks = {}
        for i in xrange(90,10,-10):
            self.info("#3 diag: sample %d%%" % i)
            sample_peaks[i]=self.__diag_peakfinding_w_control_sample(float(i)/(i+10))
        return self.__overlap (self.final_peaks, sample_peaks,top=90,bottom=10,step=-10)

    def __diag_peakfinding_w_control_sample (self, percent):
        self.treat.sample(percent) # because sampling is after
                                   # shifting, track.total is used
                                   # now.
        self.control.sample(percent)
        ratio_treat2control = float(self.treat.total)/self.control.total
        self.debug("#3 diag: after shift and merging, treat: %d, control: %d" % (self.treat.total,self.control.total))
        self.info("#3 diag: call peak candidates")
        peak_candidates = self.__call_peaks_from_trackI (self.treat)

        self.info("#3 diag: call negative peak candidates")
        negative_peak_candidates = self.__call_peaks_from_trackI (self.control)
        
        self.info("#3 diag: use control data to filter peak candidates...")
        final_peaks_percent = self.__filter_w_control(peak_candidates,self.control, ratio_treat2control)
        return final_peaks_percent
        
    def __diag_wo_control (self):
        # sample
        sample_peaks = {}
        for i in xrange(90,10,-10):
            self.info("#3 diag: sample %d%%" % i)
            sample_peaks[i]=self.__diag_peakfinding_wo_control_sample(float(i)/(i+10))
        return self.__overlap (self.final_peaks, sample_peaks,top=90,bottom=10,step=-10)

    def __diag_peakfinding_wo_control_sample (self, percent):
        self.treat.sample(percent)
        self.debug("#3 diag: after shift and merging, tags: %d" % (self.treat.total))
        self.info("#3 diag: call peak candidates")
        peak_candidates = self.__call_peaks_from_trackI (self.treat)
        self.info("#3 diag: use self to calculate local lambda and  filter peak candidates...")
        final_peaks_percent = self.__filter_w_control(peak_candidates,self.treat,1,pass_1k=True)
        return final_peaks_percent

    def __overlap (self, gold_peaks, sample_peaks, top=90,bottom=10,step=-10):
        """Calculate the overlap between several fc range for the
        golden peaks set and results from sampled data.
        
        """
        gp = PeakIO()
        gp.init_from_dict(gold_peaks)
        maxfc = int(gp.max_fold_enrichment()+1)
        diag_result = []
        for f in xrange(0,maxfc,FCSTEP):
            fc_low = f
            fc_up = f+FCSTEP
            self.debug("#3 diag: fe range = %d -- %d" % (fc_low, fc_up))
            r = self.__overlap_fc(gold_peaks, sample_peaks, fc_low, fc_up, top, bottom, step)
            if r:
                diag_result.append(r)
        return diag_result

    def __overlap_fc (self, gold_peaks, sample_peaks, fc_low, fc_up, top, bottom, step):
        ret = ["%d-%d" % (fc_low,fc_up)]
        gp = PeakIO()
        gp.init_from_dict(gold_peaks)
        gp.filter_fc(fc_low,fc_up)
        gptotal =  gp.total()
        if gptotal <= 0:
            return None

        ret.append(gptotal)
        for i in xrange(top,bottom,step):
            p = PeakIO()
            p.init_from_dict(sample_peaks[i])
            percent = 100.0*gp.overlap_with_other_peaks(p)/gptotal
            ret.append(percent)
            del p
        return ret


    def __remove_overlapping_peaks (self, peaks ):
        """peak_candidates[chrom] = [(peak_start,peak_end,peak_length,peak_summit,peak_height,number_cpr_tags)...]

        """
        new_peaks = {}
        chrs = peaks.keys()
        chrs.sort()
        for chrom in chrs:
            new_peaks[chrom]=[]
            n_append = new_peaks[chrom].append
            prev_peak = None
            peaks_chr = peaks[chrom]
            for i in xrange(len(peaks_chr)):
                if not prev_peak:
                    prev_peak = peaks_chr[i]
                    continue
                else:
                    if peaks_chr[i][0] <= prev_peak[1]:
                        s_new_peak = prev_peak[0]
                        e_new_peak = peaks_chr[i][1]
                        l_new_peak = e_new_peak-s_new_peak
                        if peaks_chr[i][4] > prev_peak[4]:
                            summit_new_peak = peaks_chr[i][3]
                            h_new_peak = peaks_chr[i][4]
                        else:
                            summit_new_peak = prev_peak[3]
                            h_new_peak = prev_peak[4]
                        prev_peak = (s_new_peak,e_new_peak,l_new_peak,summit_new_peak,h_new_peak,peaks_chr[i][5]+prev_peak[5])
                    else:
                        n_append(prev_peak)
                        prev_peak = peaks_chr[i]
            if prev_peak:
                n_append(prev_peak)
        return new_peaks
